{"format": "graph-model", "generatedBy": "2.12.0", "convertedBy": "TensorFlow.js Converter v4.5.0", "signature": {"inputs": {"xs_0": {"name": "xs_0:0", "dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "-1"}, {"size": "-1"}, {"size": "-1"}, {"size": "3"}]}}}, "outputs": {"output_0": {"name": "Identity:0", "dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "-1"}, {"size": "-1"}, {"size": "-1"}, {"size": "3"}]}}}}, "modelTopology": {"node": [{"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4787", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4783", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4785", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4747", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4741", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4739", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4737", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4735", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4731", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4733", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_975/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4749", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_969/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4751", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_977/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4773", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4767", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4765", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4763", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4761", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4757", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4725", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4723", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4721", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4719", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4715", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4717", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_981/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4759", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_965/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4775", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_959/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4777", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_967/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_983/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4705", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4701", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4703", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4665", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4659", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4657", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4655", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4653", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4649", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4651", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_949/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4667", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_943/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4669", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_951/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4691", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4685", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4683", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4681", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4679", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4675", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4643", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4641", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4639", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4637", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4633", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4635", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_955/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4677", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_939/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4693", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_933/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4695", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_941/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_957/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4511", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4507", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4509", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4471", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4465", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4463", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4461", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4459", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4455", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4457", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1051/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4473", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1045/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4475", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1053/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4497", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4491", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4489", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4487", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4485", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4481", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4449", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4447", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4445", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4443", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4439", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4441", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1057/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4483", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1041/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4499", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1035/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4501", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1043/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1059/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4429", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4425", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4427", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4389", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4383", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4381", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4379", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4377", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4373", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4375", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1025/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4391", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1019/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4393", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1027/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4415", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4409", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4407", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4405", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4403", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4399", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4367", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4365", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4363", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4361", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4357", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4359", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1031/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4401", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1015/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4417", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1009/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4419", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1017/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1033/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4235", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4231", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4233", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4195", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4189", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4187", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4185", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4183", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4179", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4181", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1127/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4197", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1121/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4199", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1129/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4221", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4215", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4213", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4211", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4209", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4205", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4173", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4171", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4169", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4167", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4163", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4165", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1133/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4207", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1117/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4223", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1111/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4225", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1119/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1135/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4153", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4149", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4151", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4113", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4107", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4105", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4103", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4101", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4097", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4099", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1101/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4115", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1095/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4117", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1103/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4139", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4133", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4131", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4129", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4127", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4123", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4091", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4089", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4087", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4085", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4081", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4083", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1107/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4125", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1091/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4141", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1085/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4143", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1093/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1109/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3989", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3987", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3979", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3969", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3973", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3967", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4017", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4015", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4007", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3997", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4001", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3995", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4045", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4043", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4035", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4025", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4029", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4023", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4073", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4071", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4063", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4053", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4057", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4051", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "512"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4265", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4263", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4255", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4245", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4249", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4243", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4293", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4291", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4283", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4273", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4277", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4271", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4321", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4319", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4311", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4301", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4305", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4299", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4349", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4347", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4339", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4329", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4333", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4327", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "256"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4541", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4539", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4531", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4521", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4525", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4519", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_911/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3611", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3609", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3607", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3605", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3601", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3603", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_917/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_913/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_919/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_925/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4569", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4567", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4559", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4549", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4553", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4547", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4597", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4595", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4587", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4577", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4581", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4575", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_873/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3787", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3785", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3783", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3781", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3777", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3779", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_879/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_875/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_881/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_887/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4625", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4623", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4615", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4605", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4609", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_4603", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_835/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3963", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3961", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3959", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3957", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3953", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3955", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_841/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_837/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_843/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_849/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3947", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3945", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3943", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3941", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3937", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3939", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_847/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_825/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3931", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3929", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3927", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3925", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3921", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3915", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3913", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3911", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3909", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3905", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3907", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_845/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3923", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_831/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_827/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_833/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_851/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3771", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3769", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3767", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3765", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3761", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3763", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_885/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_863/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3755", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3753", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3751", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3749", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3745", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3739", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3737", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3735", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3733", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3729", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3637", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3635", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3627", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3617", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3621", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3615", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3665", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3663", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3655", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3645", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3649", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3643", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3693", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3691", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3683", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3673", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3677", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3671", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3721", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3719", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3711", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3701", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3705", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3699", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "192"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3731", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_883/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3747", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_869/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_865/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_871/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_889/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3595", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3593", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3591", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3589", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3585", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3587", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_923/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_901/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3579", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3577", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3575", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3573", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3569", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3563", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3561", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3559", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3557", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3553", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3461", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3459", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3451", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3441", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3445", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3439", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3489", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3487", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3479", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3469", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3473", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3467", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3517", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3515", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3507", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3497", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3501", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3495", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3545", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3543", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3535", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3525", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3529", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3523", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "384"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3555", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_921/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3571", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_907/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_903/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_909/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_927/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3433", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3431", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3429", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3427", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3423", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3387", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3381", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3379", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3377", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3375", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3371", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3373", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1609/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3389", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1603/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3391", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1611/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3413", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3407", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3405", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3403", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3401", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3397", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3365", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3363", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3361", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3359", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3355", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3357", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1615/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3399", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1599/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3415", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1593/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3417", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1601/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1617/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3347", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3345", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3343", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3341", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3337", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3301", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3295", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3293", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3291", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3289", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3285", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3287", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1583/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3303", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1577/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3305", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1585/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3327", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3321", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3319", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3317", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3315", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3311", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3279", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3277", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3275", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3273", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3269", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3271", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1589/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3313", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1573/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3329", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1567/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3331", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1575/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1591/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3339", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1625/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1627/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3425", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1635/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1637/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3261", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3259", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3257", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3255", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3251", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3215", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3209", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3207", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3205", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3203", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3199", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3201", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1535/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3217", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1529/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3219", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1537/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3241", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3235", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3233", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3231", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3229", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3225", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3193", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3191", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3189", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3187", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3183", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3185", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1541/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3227", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1525/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3243", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1519/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3245", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1527/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1543/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3175", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3173", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3171", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3169", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3165", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3129", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3123", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3121", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3119", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3117", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3113", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3115", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1509/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3131", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1503/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3133", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1511/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3155", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3149", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3147", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3145", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3143", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3139", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3107", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3105", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3103", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3101", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3097", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3099", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1515/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3141", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1499/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3157", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1493/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3159", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1501/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1517/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3167", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1551/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1553/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3253", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1561/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1563/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "4"}, {"size": "4"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3813", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3811", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3803", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3793", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3797", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3791", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3841", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3839", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3831", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3821", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3825", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3819", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3869", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3867", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3859", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3849", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3853", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3847", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1475/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3091", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3089", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3087", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3085", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3081", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3083", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1481/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1477/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1483/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1489/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3075", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3073", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3071", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3069", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3065", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3067", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1487/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1465/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3059", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3057", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3055", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3053", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3049", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3043", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3041", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3039", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3037", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3033", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3035", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1485/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3051", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1471/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1467/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1473/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1491/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3023", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3019", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3021", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2983", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2977", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2975", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2973", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2971", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2967", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2969", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1427/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2985", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1421/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2987", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1429/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3009", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3003", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3001", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2999", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2997", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2993", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2961", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2959", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2957", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2955", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2951", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2953", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1433/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2995", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1417/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3011", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1411/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3013", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1419/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1435/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2941", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2937", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2939", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2901", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2895", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2893", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2891", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2889", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2885", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2887", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1401/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2903", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1395/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2905", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1403/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2927", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2921", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2919", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2917", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2915", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2911", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2879", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2877", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2875", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2873", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2869", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2871", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1407/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2913", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1391/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2929", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1385/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2931", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1393/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1409/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1363/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2863", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2861", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2859", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2857", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2853", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2855", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1369/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1365/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1371/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1377/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2847", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2845", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2843", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2841", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2837", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2839", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1375/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1353/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2831", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2829", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2827", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2825", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2821", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2815", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2813", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2811", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2809", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2805", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2807", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1373/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2823", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1359/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1355/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1361/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1379/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2795", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2791", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2793", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2755", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2749", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2747", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2745", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2743", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2739", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2741", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1315/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2757", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1309/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2759", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1317/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2781", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2775", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2773", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2771", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2769", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2765", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2733", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2731", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2729", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2727", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2723", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2725", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1321/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2767", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1305/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2783", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1299/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2785", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1307/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1323/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2713", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2709", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2711", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2673", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2667", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2665", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2663", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2661", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2657", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2659", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1289/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2675", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1283/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2677", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1291/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2699", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2693", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2691", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2689", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2687", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2683", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2651", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2649", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2647", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2645", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2641", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2643", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1295/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2685", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1279/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2701", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1273/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2703", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1281/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1297/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1251/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2635", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2633", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2631", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2629", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2625", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2627", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1257/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1253/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1259/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1265/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2619", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2617", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2615", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2613", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2609", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2611", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1263/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1241/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2603", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2601", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2599", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2597", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2593", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2587", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2585", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2583", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2581", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2577", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2579", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1261/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2595", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1247/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1243/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1249/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1267/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2567", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2563", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2565", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2527", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2521", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2519", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2517", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2515", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2511", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2513", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1203/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2529", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1197/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2531", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1205/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2553", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2547", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2545", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2543", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2541", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2537", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2505", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2503", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2501", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2499", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2495", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2497", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1209/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2539", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1193/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2555", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1187/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2557", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1195/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1211/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2485", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2481", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2483", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2445", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2439", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2437", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2435", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2433", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2429", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2431", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1177/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2447", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1171/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2449", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1179/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2471", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2465", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2463", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2461", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2459", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2455", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2423", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2421", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2419", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2417", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2413", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2415", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1183/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2457", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1167/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2473", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1161/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2475", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1169/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1185/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1665/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2279", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2277", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2275", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2273", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2269", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2271", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1671/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1667/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1673/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1679/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2263", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2261", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2259", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2257", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2253", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2255", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1677/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1655/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2247", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2245", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2243", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2241", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2237", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2231", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2229", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2227", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2225", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2221", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2223", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1675/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2239", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1661/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1657/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1663/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1681/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2211", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2207", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2209", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2171", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2165", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2163", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2161", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2159", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2155", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2157", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_233/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2173", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_227/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2175", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_235/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2197", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2191", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2189", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2187", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2185", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2181", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2149", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2147", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2145", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2143", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2139", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2141", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_239/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2183", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_223/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2199", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_217/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2201", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_225/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_241/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2129", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2125", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2127", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2089", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2083", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2081", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2079", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2077", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2073", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2075", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_207/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2091", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_201/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2093", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_209/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2115", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2109", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2107", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2105", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2103", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2099", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2067", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2065", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2063", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2061", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2057", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2059", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_213/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2101", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_197/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2117", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_191/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2119", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_199/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_215/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1965", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1963", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1955", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1945", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1949", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1943", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1993", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1991", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1983", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1973", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1977", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1971", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2021", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2019", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2011", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2001", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2005", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1999", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2049", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2047", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2039", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2029", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2033", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2027", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "4"}, {"size": "4"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1701/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2343", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2341", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2339", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2337", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2333", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2335", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1707/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1703/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1709/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1715/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2327", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2325", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2323", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2321", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2317", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2319", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1713/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1691/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2311", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2309", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2307", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2305", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2301", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2295", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2293", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2291", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2289", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2285", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2287", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1711/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2303", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1697/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1693/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1699/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1717/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1935", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1931", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1933", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1895", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1889", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1887", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1885", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1883", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1879", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1881", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_309/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1897", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_303/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1899", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_311/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1921", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1915", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1913", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1911", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1909", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1905", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1873", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1871", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1869", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1867", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1863", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1865", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_315/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1907", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_299/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1923", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_293/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1925", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_301/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_317/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1853", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1849", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1851", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1813", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1807", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1805", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1803", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1801", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1797", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1799", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_283/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1815", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_277/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1817", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_285/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1839", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1833", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1831", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1829", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1827", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1823", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1791", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1789", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1787", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1785", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1781", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1783", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_289/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1825", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_273/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1841", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_267/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1843", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_275/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_291/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1689", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1687", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1679", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1669", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1673", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1667", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1717", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1715", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1707", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1697", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1701", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1695", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1745", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1743", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1735", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1725", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1729", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1723", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1773", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1771", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1763", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1753", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1757", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1751", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "256"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "96"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "4"}, {"size": "4"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1737/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2407", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2405", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2403", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2401", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2397", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2399", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1743/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1739/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1745/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1751/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2391", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2389", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2387", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2385", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2381", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2383", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1749/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1727/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2375", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2373", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2371", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2369", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2365", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2359", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2357", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2355", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2353", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2349", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2351", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1747/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_2367", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1733/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1729/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1735/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_1753/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1659", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1655", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1657", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1619", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1613", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1611", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1609", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1607", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1603", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1605", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_385/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1621", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_379/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1623", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_387/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1645", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1639", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1637", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1635", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1633", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1629", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1597", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1595", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1593", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1591", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1587", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1589", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_391/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1631", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_375/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1647", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_369/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1649", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_377/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_393/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1577", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1573", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1575", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1537", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1531", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1529", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1527", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1525", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1521", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1523", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_359/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1539", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_353/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1541", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_361/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1563", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1557", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1555", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1553", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1551", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1547", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1515", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1513", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1511", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1509", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1505", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1507", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_365/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1549", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_349/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1565", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_343/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1567", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_351/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_367/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_169/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1035", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1033", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1031", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1029", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1025", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1027", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_175/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_171/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_177/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_183/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1413", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1411", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1403", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1393", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1397", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1391", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_131/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1211", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1209", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1207", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1205", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1201", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1203", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_137/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_133/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_139/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_145/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1441", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1439", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1431", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1421", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1425", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1419", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1469", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1467", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1459", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1449", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1453", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1447", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_93/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1387", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1385", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1383", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1381", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1377", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1379", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_99/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_95/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_101/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_107/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1371", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1369", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1367", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1365", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1361", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1363", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_105/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_83/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1355", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1353", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1351", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1349", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1345", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1339", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1337", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1335", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1333", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1329", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1331", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_103/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1347", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_89/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_85/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_91/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_109/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1195", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1193", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1191", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1189", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1185", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1187", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_143/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_121/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1179", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1177", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1175", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1173", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1169", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1163", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1161", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1159", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1157", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1153", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1061", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1059", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1051", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1041", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1045", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1039", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1089", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1087", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1079", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1069", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1073", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1067", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1117", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1115", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1107", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1097", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1101", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1095", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1145", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1143", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1135", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1125", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1129", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1123", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "192"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1155", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_141/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1171", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_127/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_123/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_129/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_147/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1019", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1017", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1015", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1013", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1009", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1011", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_181/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_159/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1003", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1001", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_999", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_997", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_993", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_987", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_985", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_983", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_981", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_977", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_885", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_883", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_875", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_865", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_869", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_863", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_913", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_911", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_903", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_893", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_897", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_891", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_941", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_939", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_931", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_921", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_925", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_919", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_969", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_967", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_959", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_949", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_953", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_947", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "384"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_979", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_179/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_995", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_165/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_161/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_167/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_185/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/transpose_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "2"}, {"size": "2"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_857", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_855", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_853", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_851", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_847", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_811", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_805", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_803", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_801", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_799", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_795", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_797", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_759/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_813", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_753/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_815", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_761/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_837", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_831", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_829", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_827", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_825", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_821", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_789", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_787", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_785", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_783", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_779", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_781", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_765/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_823", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_749/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_839", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_743/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_841", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_751/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_767/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_771", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_769", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_767", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_765", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_761", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_725", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_719", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_717", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_715", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_713", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_709", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_711", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_733/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_727", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_727/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_729", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_735/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_751", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_745", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_743", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_741", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_739", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_735", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_703", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_701", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_699", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_697", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_693", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_695", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_739/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_737", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_723/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_753", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_717/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_755", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_725/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_741/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_763", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_775/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_777/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_849", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_785/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_787/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_685", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_683", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_681", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_679", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_675", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_639", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_633", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_631", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_629", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_627", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_623", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_625", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_685/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_641", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_679/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_643", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_687/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_665", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_659", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_657", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_655", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_653", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_649", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_617", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_615", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_613", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_611", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_607", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_609", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_691/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_651", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_675/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_667", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_669/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_669", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_677/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_693/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_599", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_597", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_595", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_593", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_589", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_553", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_547", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_545", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_543", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_541", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_537", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_539", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_659/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_555", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_653/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_557", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_661/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_579", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_573", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_571", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_569", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_567", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_563", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_531", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_529", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_527", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_525", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_521", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_523", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_665/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_565", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_649/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_581", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_643/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_583", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_651/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_667/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_591", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_701/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_703/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_677", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_711/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_713/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "4"}, {"size": "4"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1237", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1235", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1227", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1217", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1221", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1215", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1265", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1263", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1255", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1245", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1249", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1243", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1293", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1291", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1283", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1273", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1277", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1271", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_511", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_507", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_509", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_471", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_465", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_463", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_461", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_459", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_455", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_457", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_613/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_473", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_607/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_475", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_615/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_497", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_491", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_489", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_487", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_485", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_481", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_449", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_447", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_445", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_443", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_439", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_441", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_619/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_483", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_603/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_499", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_597/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_501", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_605/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_621/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_429", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_425", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_427", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_389", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_383", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_381", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_379", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_377", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_373", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_375", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_587/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_391", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_581/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_393", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_589/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_415", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_409", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_407", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_405", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_403", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_399", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_367", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_365", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_363", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_361", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_357", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_359", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_593/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_401", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_577/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_417", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_571/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_419", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_579/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_595/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_347", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_343", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_345", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_307", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_301", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_299", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_297", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_295", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_291", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_293", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_537/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_309", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_531/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_311", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_539/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_333", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_327", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_325", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_323", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_321", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_317", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_285", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_283", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_281", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_279", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_275", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_277", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_543/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_319", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_527/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_335", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_521/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_337", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_529/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_545/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_265", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_261", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_263", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "16"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "16"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_225", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_219", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_217", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_215", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_213", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_209", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_211", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_511/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_227", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_505/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_229", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_513/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_251", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_245", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_243", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_241", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_239", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_235", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_203", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_201", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_199", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_197", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_193", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_195", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_517/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_237", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_501/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_253", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_495/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_255", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_503/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_519/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "128"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_183", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_179", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_181", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_143", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_137", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_135", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_133", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_131", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_127", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_129", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_461/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_145", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_455/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_147", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_463/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_169", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_163", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_161", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_159", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_157", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_153", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_121", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_119", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_117", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_115", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_111", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_113", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_467/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_155", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_451/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_171", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_445/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_173", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_453/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_469/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_101", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_97", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_99", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "2"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "8"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "8"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_53", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_435/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_429/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_437/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_87", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_81", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_79", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_77", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_75", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_37", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_441/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_425/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_89", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_419/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "256"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "256"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_91", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_427/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "6"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_arg_443/ReadVariableOp", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "64"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "4"}, {"size": "4"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_11", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_25", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_29", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_31", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_15", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_35", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_37", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_53", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "96"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "4"}, {"size": "4"}, {"size": "64"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "64"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "192"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1321", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1319", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1311", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1301", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1305", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1299", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "96"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1497", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1495", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1487", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1477", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1481", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1475", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "512"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_19", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_75", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_69", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_17", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_71", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_73", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_77", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_81", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_85", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_91", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_93", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_97", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_99", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_23", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_109", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_103", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_21", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_105", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_107", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_111", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_115", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_121", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_125", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_127", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Sign_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/add/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_131", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_133", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape/3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "3"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "192"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "128"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3897", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3895", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_65", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3887", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_6", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_39", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3877", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3881", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_41", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_45", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/start", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/delta", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_51", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_55", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape/1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_3875", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_9", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Sign_4", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/add_1/y", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_1", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_2", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "1"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/Const_5", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose/perm", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {"dim": [{"size": "4"}]}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "96"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "32"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat/axis", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "128"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "32"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/transpose_3", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "3"}, {"size": "3"}, {"size": "32"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Reshape", "op": "Const", "attr": {"value": {"tensor": {"dtype": "DT_FLOAT", "tensorShape": {"dim": [{"size": "1"}, {"size": "1"}, {"size": "1"}, {"size": "3"}]}}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "xs_0", "op": "Placeholder", "attr": {"shape": {"shape": {"dim": [{"size": "-1"}, {"size": "-1"}, {"size": "-1"}, {"size": "3"}]}}, "dtype": {"type": "DT_FLOAT"}}, "inputs": []}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Conv2D", "op": "Conv2D", "input": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Conv2D", "op": "Conv2D", "input": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Shape", "op": "Shape", "input": ["xs_0"], "attr": {"T": {"type": "DT_FLOAT"}, "out_type": {"type": "DT_INT32"}}, "experimentalType": {"typeId": "TFT_PRODUCT", "args": [{"typeId": "TFT_SHAPE_TENSOR", "args": [{"typeId": "TFT_INT32"}]}]}, "inputs": ["xs_0"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Shape_1", "op": "Shape", "input": ["xs_0"], "attr": {"T": {"type": "DT_FLOAT"}, "out_type": {"type": "DT_INT32"}}, "experimentalType": {"typeId": "TFT_PRODUCT", "args": [{"typeId": "TFT_SHAPE_TENSOR", "args": [{"typeId": "TFT_INT32"}]}]}, "inputs": ["xs_0"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Shape_2", "op": "Shape", "input": ["xs_0"], "attr": {"T": {"type": "DT_FLOAT"}, "out_type": {"type": "DT_INT32"}}, "experimentalType": {"typeId": "TFT_PRODUCT", "args": [{"typeId": "TFT_SHAPE_TENSOR", "args": [{"typeId": "TFT_INT32"}]}]}, "inputs": ["xs_0"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Shape", "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "1"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_INT32"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Shape", "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Shape_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "1"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_INT32"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Shape_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Shape_2", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_2"], "attr": {"shrink_axis_mask": {"i": "1"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_INT32"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Shape_2", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_1", "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Sign", "op": "Sign", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/FloorMod", "op": "FloorMod", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_1", "StatefulPartitionedCall/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_1", "StatefulPartitionedCall/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/floordiv", "op": "FloorDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_1", "StatefulPartitionedCall/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_1", "StatefulPartitionedCall/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Sign_3", "op": "Sign", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_2"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/FloorMod_2", "op": "FloorMod", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_2", "StatefulPartitionedCall/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_2", "StatefulPartitionedCall/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/floordiv_1", "op": "FloorDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_2", "StatefulPartitionedCall/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice_2", "StatefulPartitionedCall/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Sign", "StatefulPartitionedCall/jax2tf_apply_/Sign_1"], "attr": {"T": {"type": "DT_INT32"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Sign", "StatefulPartitionedCall/jax2tf_apply_/Sign_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/FloorMod", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1/y"], "attr": {"T": {"type": "DT_INT32"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/FloorMod", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/add", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/floordiv", "StatefulPartitionedCall/jax2tf_apply_/add/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/floordiv", "StatefulPartitionedCall/jax2tf_apply_/add/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_2", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Sign_3", "StatefulPartitionedCall/jax2tf_apply_/Sign_4"], "attr": {"T": {"type": "DT_INT32"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Sign_3", "StatefulPartitionedCall/jax2tf_apply_/Sign_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/FloorMod_2", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3/y"], "attr": {"T": {"type": "DT_INT32"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/FloorMod_2", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/add_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/floordiv_1", "StatefulPartitionedCall/jax2tf_apply_/add_1/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/floordiv_1", "StatefulPartitionedCall/jax2tf_apply_/add_1/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/NotEqual_2", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/NotEqual_2", "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/add", "StatefulPartitionedCall/jax2tf_apply_/floordiv"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/add", "StatefulPartitionedCall/jax2tf_apply_/floordiv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/add_1", "StatefulPartitionedCall/jax2tf_apply_/floordiv_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/add_1", "StatefulPartitionedCall/jax2tf_apply_/floordiv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_21"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_23"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_30", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_75"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_75"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_69"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_69"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_27", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_71"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_71"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_29", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_31", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_77"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_77"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_32", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_81"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_81"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_33", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_85"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_85"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_35", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_91"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_91"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_37", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_97"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_97"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_48", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_125"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_125"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_50", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_131"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_131"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "StatefulPartitionedCall/jax2tf_apply_/SelectV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_10", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_25"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_12", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_31"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_35"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_35"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_14", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_37"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_37"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_53"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_53"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_36", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_93"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_93"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_38", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_99"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_99"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_43", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_109"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_109"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_103"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_103"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_40", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_105"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_105"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_42", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_107"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_107"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_44", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_111"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_111"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_45", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_115"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_115"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_47", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_121"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_121"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_49", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_127"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_127"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_51", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_133"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_133"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_65"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_39"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_41"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_41"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_45"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_45"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_21", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_51"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_51"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_22", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_55"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_55"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_9"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_13"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_12"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_14"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_7"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_8"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_12"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape/1"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_7"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_27"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_27"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_29"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_29"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_31", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_31", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_32", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_32", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape/1"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_33"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_33"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_14", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_17", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"], "attr": {"N": {"i": "6"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_12", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_14"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_14"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_16"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_16"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape/1"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_21"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_23", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_25", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_25", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_35", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_36", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_35", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_36", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_37", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_38", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_37", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_38", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_40"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_40"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_42"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_42"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_44", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_44", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_45", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_45", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape/1"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_47"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_48", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_49", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_48", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_49", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_50", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_51", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_50", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_51", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_23"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_24"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_24"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_16", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_18"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_18"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3", "op": "Range", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/delta"], "attr": {"Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/start", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/delta"]}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape_const_axis", "op": "Const", "input": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20"], "attr": {"value": {"tensor": {"dtype": "DT_INT32", "tensorShape": {}}}, "dtype": {"type": "DT_INT32"}}, "inputs": ["^StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_21", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape/1"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_22", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape/1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_18", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"], "attr": {"N": {"i": "2"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_24", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/jit_fn_/Mul_25"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2221", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2221", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3987"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3987"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4043"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4043"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4263"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4263"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4319"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4319"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4539"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4539"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4595"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_4595"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3635"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3635"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3691"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3691"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3459"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3459"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3515"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3515"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3811"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3811"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3867"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_3867"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1963"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1963"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_2019"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_2019"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1687"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1687"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1743"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1743"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1411"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1411"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1467"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1467"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1059"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1059"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1115"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1115"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_883"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_883"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_939"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_939"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1235"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1235"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_4", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1291"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_3", "StatefulPartitionedCall/Const_1291"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_4", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_9", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_7", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_7", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_26", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_33", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_33", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4015"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4015"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4071"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4071"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4291"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4291"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4347"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4347"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4567"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4567"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4623"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_4623"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3663"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3663"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3719"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3719"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3487"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3487"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3543"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3543"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3839"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3839"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1991"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1991"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_2047"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_2047"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1715"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1715"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1771"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1771"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1439"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1439"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1087"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1087"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1143"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1143"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_911"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_911"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_967"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_967"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_19", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1263"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1263"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape", "op": "Pack", "input": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"], "attr": {"N": {"i": "4"}, "axis": {"i": "0"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_13", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_21", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_21", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1319"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1319"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1495"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_1495"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_39", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_47", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_47", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_9", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3895"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_7", "StatefulPartitionedCall/Const_3895"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_15", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape", "op": "ExpandDims", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape_const_axis"], "attr": {"Tdim": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_20", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape_const_axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_29", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_29", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_2", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_5", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2223"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2223"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3967"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3967"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4023"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4023"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4243"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4243"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4299"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4299"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4519"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4519"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4575"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_4575"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3615"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3615"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3671"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3671"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3439"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3439"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3495"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3495"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3791"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3791"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3847"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_3847"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1943"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1943"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1999"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1999"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1667"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1667"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1723"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1723"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1391"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1391"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1447"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1447"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1039"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1039"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1095"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1095"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_863"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_863"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_919"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_919"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1215"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1215"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1271"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_2", "StatefulPartitionedCall/Const_1271"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_4", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3995"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3995"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4051"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4051"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4271"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4271"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4327"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4327"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4547"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4547"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4603"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_4603"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3643"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3643"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3699"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3699"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3467"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3467"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3523"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3523"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3819"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3819"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1971"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1971"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_2027"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_2027"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1695"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1695"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1751"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1751"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1419"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1419"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1067"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1067"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1123"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1123"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_891"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_891"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_947"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_947"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1243"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1243"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1299"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1299"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1475"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_1475"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_6", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_2", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7"], "attr": {"SrcT": {"type": "DT_INT32"}, "DstT": {"type": "DT_FLOAT"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3875"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_5", "StatefulPartitionedCall/Const_3875"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_31"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_31"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast", "StatefulPartitionedCall/Const_3969"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast", "StatefulPartitionedCall/Const_3969"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast", "StatefulPartitionedCall/Const_4025"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast", "StatefulPartitionedCall/Const_4025"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast", "StatefulPartitionedCall/Const_4245"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast", "StatefulPartitionedCall/Const_4245"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast", "StatefulPartitionedCall/Const_4301"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast", "StatefulPartitionedCall/Const_4301"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast", "StatefulPartitionedCall/Const_4521"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast", "StatefulPartitionedCall/Const_4521"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast", "StatefulPartitionedCall/Const_4577"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast", "StatefulPartitionedCall/Const_4577"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast", "StatefulPartitionedCall/Const_3617"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast", "StatefulPartitionedCall/Const_3617"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast", "StatefulPartitionedCall/Const_3673"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast", "StatefulPartitionedCall/Const_3673"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast", "StatefulPartitionedCall/Const_3441"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast", "StatefulPartitionedCall/Const_3441"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast", "StatefulPartitionedCall/Const_3497"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast", "StatefulPartitionedCall/Const_3497"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast", "StatefulPartitionedCall/Const_3793"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast", "StatefulPartitionedCall/Const_3793"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast", "StatefulPartitionedCall/Const_3849"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast", "StatefulPartitionedCall/Const_3849"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast", "StatefulPartitionedCall/Const_1945"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast", "StatefulPartitionedCall/Const_1945"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast", "StatefulPartitionedCall/Const_2001"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast", "StatefulPartitionedCall/Const_2001"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast", "StatefulPartitionedCall/Const_1669"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast", "StatefulPartitionedCall/Const_1669"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast", "StatefulPartitionedCall/Const_1725"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast", "StatefulPartitionedCall/Const_1725"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast", "StatefulPartitionedCall/Const_1393"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast", "StatefulPartitionedCall/Const_1393"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast", "StatefulPartitionedCall/Const_1449"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast", "StatefulPartitionedCall/Const_1449"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast", "StatefulPartitionedCall/Const_1041"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast", "StatefulPartitionedCall/Const_1041"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast", "StatefulPartitionedCall/Const_1097"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast", "StatefulPartitionedCall/Const_1097"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast", "StatefulPartitionedCall/Const_865"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast", "StatefulPartitionedCall/Const_865"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast", "StatefulPartitionedCall/Const_921"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast", "StatefulPartitionedCall/Const_921"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast", "StatefulPartitionedCall/Const_1217"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast", "StatefulPartitionedCall/Const_1217"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast", "StatefulPartitionedCall/Const_1273"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast", "StatefulPartitionedCall/Const_1273"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast", "StatefulPartitionedCall/Const_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast", "StatefulPartitionedCall/Const_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_4", "StatefulPartitionedCall/Const_17"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_4", "StatefulPartitionedCall/Const_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3997"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3997"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4053"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4053"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4273"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4273"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4329"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4329"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4549"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4549"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4605"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_4605"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3645"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3645"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3701"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3701"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3469"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3469"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3525"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3525"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3821"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3821"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1973"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1973"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_2029"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_2029"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1697"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1697"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1753"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1753"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1421"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1421"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1069"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1069"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1125"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1125"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_893"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_893"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_949"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_949"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1245"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1245"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_13"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_13"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1301"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1301"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1477"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_1477"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_6", "StatefulPartitionedCall/Const_21"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_6", "StatefulPartitionedCall/Const_21"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3877"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_2", "StatefulPartitionedCall/Const_3877"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_28", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_15", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_41", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/AddV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_16", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3973"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3973"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4029"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4029"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4249"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4249"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4305"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4305"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4525"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4525"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4581"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_4581"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3621"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3621"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3677"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3677"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3445"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3445"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3501"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3501"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3797"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3797"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3853"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_3853"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1949"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1949"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_2005"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_2005"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1673"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1673"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1729"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1729"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1397"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1397"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1453"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1453"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1045"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1045"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1101"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1101"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_869"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_869"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_925"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_925"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1221"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1221"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1277"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_5", "StatefulPartitionedCall/Const_1277"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_2", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_28", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_28", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4001"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4001"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4057"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4057"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4277"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4277"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4333"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4333"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4553"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4553"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4609"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_4609"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3649"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3649"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3705"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3705"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3473"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3473"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3529"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3529"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3825"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3825"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1977"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1977"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_2033"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_2033"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1701"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1701"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1757"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1757"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1425"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1425"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1073"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1073"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1129"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1129"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_897"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_897"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_953"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_953"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1249"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1249"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_1", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_15", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1305"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1305"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1481"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_1481"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_41", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_41", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3881"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Mul_16", "StatefulPartitionedCall/Const_3881"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3989"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3989"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4045"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4045"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4265"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4265"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4321"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4321"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4541"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4541"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4597"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_4597"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3637"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3637"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3693"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3693"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3461"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3461"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3517"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3517"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3813"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3813"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3869"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_3869"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1965"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1965"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_2021"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_2021"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1689"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1689"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1745"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1745"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1413"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1413"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1469"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1469"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1061"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1061"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1117"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1117"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_885"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_885"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_941"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_941"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1237"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1237"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1293"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/Const_1293"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/LessEqual", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_2", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4017"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4017"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4073"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4073"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4293"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4293"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4349"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4349"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4569"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4569"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4625"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_4625"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3665"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3665"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3721"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3721"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3489"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3489"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3545"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3545"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3841"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3841"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1993"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1993"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_2049"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_2049"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1717"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1717"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1773"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1773"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1441"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1441"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1089"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1089"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1145"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1145"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_913"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_913"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_969"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_969"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1265"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1265"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_1", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1321"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1321"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1497"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_1497"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_3", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/GreaterEqual_1", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3897"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/Const_3897"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/LessEqual_1", "op": "LessEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1675/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1675/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_op_/LogicalAnd", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/LessEqual"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/LessEqual"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign", "op": "Sign", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_2", "op": "Sign", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_1", "op": "Sign", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_3", "op": "Sign", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_op_/LogicalAnd_1", "op": "LogicalAnd", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/LessEqual_1"], "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/GreaterEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/LessEqual_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_441/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_441/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_op_/LogicalAnd", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_2", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit__lambda__/Floor_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_op_/LogicalAnd_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_7", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2231", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2231", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2225", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2225", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_2", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_1", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_3", "op": "Floor", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_3", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_39", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_39", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_33", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_33", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_3", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sign_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Floor_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_6", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_1", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_1"], "attr": {"SrcT": {"type": "DT_FLOAT"}, "DstT": {"type": "DT_INT32"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_5"], "attr": {"SrcT": {"type": "DT_FLOAT"}, "DstT": {"type": "DT_INT32"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_3"], "attr": {"SrcT": {"type": "DT_FLOAT"}, "DstT": {"type": "DT_INT32"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7", "op": "Cast", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_7"], "attr": {"SrcT": {"type": "DT_FLOAT"}, "DstT": {"type": "DT_INT32"}, "Truncate": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_4", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_6"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2227", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2227", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_3", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less", "op": "Less", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1", "StatefulPartitionedCall/Const_11"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1", "StatefulPartitionedCall/Const_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_2", "op": "Less", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5", "StatefulPartitionedCall/Const_19"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5", "StatefulPartitionedCall/Const_19"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_30"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_30"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_1", "op": "Less", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3", "StatefulPartitionedCall/Const_15"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3", "StatefulPartitionedCall/Const_15"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_17"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_17"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_3", "op": "Less", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7", "StatefulPartitionedCall/Const_23"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7", "StatefulPartitionedCall/Const_23"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_43"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Mul_43"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_8", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_35", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_35", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/jit_fn_/Less_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/AddV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Cast_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_8"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Sub_8"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2229", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2229", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/SelectV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_37", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_37", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape"], "attr": {"T": {"type": "DT_INT32"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze", "op": "Squeeze", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1"], "attr": {"T": {"type": "DT_INT32"}, "squeeze_dims": {"list": {"i": ["-1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_2", "op": "Squeeze", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7"], "attr": {"T": {"type": "DT_INT32"}, "squeeze_dims": {"list": {"i": ["-1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_1", "op": "Squeeze", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4"], "attr": {"T": {"type": "DT_INT32"}, "squeeze_dims": {"list": {"i": ["-1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_3", "op": "Squeeze", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10"], "attr": {"T": {"type": "DT_INT32"}, "squeeze_dims": {"list": {"i": ["-1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_2", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/NotEqual", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_2"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/Minimum", "op": "Minimum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/Minimum", "op": "Minimum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/Minimum", "op": "Minimum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/Minimum", "op": "Minimum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Squeeze_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_5", "op": "Abs", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/NotEqual_1", "op": "NotEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_6"], "attr": {"T": {"type": "DT_FLOAT"}, "incompatible_shape_error": {"b": true}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3979"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3979"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4035"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4035"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4255"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4255"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4311"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4311"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4531"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4531"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4587"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_4587"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3627"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3627"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3683"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3683"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3451"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3451"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3507"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3507"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3803"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3803"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3859"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_3859"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1955"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1955"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_2011"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_2011"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1679"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1679"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1735"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1735"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1403"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1403"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1459"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1459"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1051"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1051"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1107"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1107"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_875"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_875"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_931"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_931"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1227"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1227"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Greater", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1283"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_2", "StatefulPartitionedCall/Const_1283"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/NotEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4007"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4007"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4063"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4063"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4283"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4283"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4339"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4339"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4559"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4559"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4615"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_4615"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3655"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3655"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3711"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3711"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3479"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3479"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3535"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3535"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3831"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3831"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1983"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1983"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_2039"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_2039"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1707"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1707"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1763"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1763"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1431"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1431"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1079"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1079"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1135"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1135"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_903"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_903"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_959"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_959"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1255"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1255"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1311"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1311"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1487"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_1487"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/y"], "attr": {"T": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/Minimum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/y"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Greater_1", "op": "Greater", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3887"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_absolute_/Abs_5", "StatefulPartitionedCall/Const_3887"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_3", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/NotEqual_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Greater", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_4", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2", "op": "GatherV2", "input": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2/axis"], "attr": {"Taxis": {"type": "DT_INT32"}, "Tparams": {"type": "DT_FLOAT"}, "Tindices": {"type": "DT_INT32"}, "batch_dims": {"i": "0"}}, "inputs": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2", "op": "GatherV2", "input": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2/axis"], "attr": {"Tparams": {"type": "DT_FLOAT"}, "Taxis": {"type": "DT_INT32"}, "Tindices": {"type": "DT_INT32"}, "batch_dims": {"i": "0"}}, "inputs": ["xs_0", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_BOOL"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"], "attr": {"T": {"type": "DT_BOOL"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_fn_/Greater_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_7", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Maximum_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2237", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2237", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_1", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_4", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit_true_divide_/truediv_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_5", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_9"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2239"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2239"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1", "op": "GatherV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1/axis"], "attr": {"Taxis": {"type": "DT_INT32"}, "Tparams": {"type": "DT_FLOAT"}, "Tindices": {"type": "DT_INT32"}, "batch_dims": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3", "op": "GatherV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3/axis"], "attr": {"Tparams": {"type": "DT_FLOAT"}, "Taxis": {"type": "DT_INT32"}, "Tindices": {"type": "DT_INT32"}, "batch_dims": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_45", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_45", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_71", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_71", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1661/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1661/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_47"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_47"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_73"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_73"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2247", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2247", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2241", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2241", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2243", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2243", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_435/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_435/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_425/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_425/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2285", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2285", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2349", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2349", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2245", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2245", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_55", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_55", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_49", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_49", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_81", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_81", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_75", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_75", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2287"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2287"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2351"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2351"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_51", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_51", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_77", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_77", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_53", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_53", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_79", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_79", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1711/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1711/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1747/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1747/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1655/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1655/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1657/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1657/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2295", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2295", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2289", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2289", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2359", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2359", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2353", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2353", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2291", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2291", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2355", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2355", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2293", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2293", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2357", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2357", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_61", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_61", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_87", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_87", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1663/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1663/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_63"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_63"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_89"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_89"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2301", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2301", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2303"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2303"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2367"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2367"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_429/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_429/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_419/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_419/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_65"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_65"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_91"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_91"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1697/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1697/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1733/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1733/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_437/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_437/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_427/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_427/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2311", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2311", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2305", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2305", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2375", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2375", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2369", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2369", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2307", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2307", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2371", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2371", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2309", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2309", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2373", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2373", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_443/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_443/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_97", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_97", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1691/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1691/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1693/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1693/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1727/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1727/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1729/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1729/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_99"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_99"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1699/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1699/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1735/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1735/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_101", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_101", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_111", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_111", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_113"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_113"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_467/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_467/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_121", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_121", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_115", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_115", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_117", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_117", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_119", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_119", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_127", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_127", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_153", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_153", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_129"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_129"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_155"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_155"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_461/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_461/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_451/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_451/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_137", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_137", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_131", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_131", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_157", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_157", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_133", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_133", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_159", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_159", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_135", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_135", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_161", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_161", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_143", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_143", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_145"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_145"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_171"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_171"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_455/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_455/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_445/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_445/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_147"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_147"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_173"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_173"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_463/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_463/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_453/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_453/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_469/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_469/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_179", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_179", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_181"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_181"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_183", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_183", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_193", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_193", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_195"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_195"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_517/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_517/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_203", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_203", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_197", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_197", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_199", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_199", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_201", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_201", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_235", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_235", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_211"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_211"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_237"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_237"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_511/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_511/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_501/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_501/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_219", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_219", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_213", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_213", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_245", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_245", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_239", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_239", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_215", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_215", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_241", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_241", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_217", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_217", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_243", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_243", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_225", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_225", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_251", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_251", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_227"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_227"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_253"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_253"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_505/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_505/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_495/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_495/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_229"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_229"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_255"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_255"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_513/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_513/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_503/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_503/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_519/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_519/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_261", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_261", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_263"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_263"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_265", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_265", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_275", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_275", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_277"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_277"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_543/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_543/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_285", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_285", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_279", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_279", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_281", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_281", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_283", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_283", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_291", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_291", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_317", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_317", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_293"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_293"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_319"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_319"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_537/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_537/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_527/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_527/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_301", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_301", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_295", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_295", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_327", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_327", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_321", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_321", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_297", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_297", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_323", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_323", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_299", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_299", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_325", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_325", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_307", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_307", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_333", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_333", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_309"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_309"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_335"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_335"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_531/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_531/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_521/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_521/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_311"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_311"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_337"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_337"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_539/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_539/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_529/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_529/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_545/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_545/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_343", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_343", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_345"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_345"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_347", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_347", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_357", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_357", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_359"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_359"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_593/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_593/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_367", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_367", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_363", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_363", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_373", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_373", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_399", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_399", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_375"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_375"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_401"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_401"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_587/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_587/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_577/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_577/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_383", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_383", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_409", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_409", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_379", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_379", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_389", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_389", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_415", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_415", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_391"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_391"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_417"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_417"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_581/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_581/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_571/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_571/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_393"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_393"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_419"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_419"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_589/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_589/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_579/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_579/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_595/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_595/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_425", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_425", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_427"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_427"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_439", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_439", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_441"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_441"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_619/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_619/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_449", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_449", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_443", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_443", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_445", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_445", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_447", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_447", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_455", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_455", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_481", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_481", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_457"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_457"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_483"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_483"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_613/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_613/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_603/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_603/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_465", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_465", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_459", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_459", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_491", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_491", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_485", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_485", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_461", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_461", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_487", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_487", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_463", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_463", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_489", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_489", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_471", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_471", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_497", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_497", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_473"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_473"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_499"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_499"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_607/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_607/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_597/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_597/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_475"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_475"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_501"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_501"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_615/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_615/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_605/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_605/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_621/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_621/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_507", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_507", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_509"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_509"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat/axis"], "attr": {"N": {"i": "3"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1/axis"], "attr": {"N": {"i": "3"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2/axis"], "attr": {"N": {"i": "3"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_521", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_521", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_977", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_977", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_523"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_523"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_979"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_979"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1153", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1153", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1329", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1329", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_1155"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_1155"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_1331"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_1331"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_665/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_665/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_179/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_179/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_531", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_531", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_525", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_525", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_987", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_987", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_981", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_981", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_141/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_141/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_103/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_103/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_527", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_527", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1157", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1157", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_983", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_983", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1339", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1339", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1333", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1333", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_529", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_529", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_985", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_985", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1159", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1159", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1335", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1335", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1161", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1161", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1337", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1337", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_993", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_993", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_995"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_995"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1345", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1345", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_537", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_537", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1171"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1171"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1347"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1347"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_539"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_539"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_565"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_565"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_165/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_165/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1003", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1003", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_997", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_997", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_127/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_127/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_89/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_89/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_659/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_659/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_649/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_649/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1179", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1179", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1173", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1173", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_999", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_999", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1355", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1355", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1349", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1349", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_547", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_547", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_541", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_541", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_573", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_573", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_567", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_567", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1001", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1001", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1175", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1175", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1351", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1351", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_543", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_543", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_569", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_569", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1177", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1177", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1353", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1353", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_545", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_545", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_571", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_571", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_159/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_159/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_161/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_161/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_121/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_121/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_123/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_123/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_83/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_83/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_85/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_85/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_579", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_579", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_555"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_555"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_581"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_581"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_167/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_167/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_129/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_129/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_91/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_91/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_653/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_653/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_643/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_643/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_557"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_557"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_583"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_583"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_661/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_661/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_651/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_651/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_667/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_667/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_589", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_589", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_591"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_591"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_701/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_701/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_599", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_599", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_595", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_595", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_597", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_597", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_703/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_703/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_607", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_607", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_609"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_609"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_691/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_691/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_617", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_617", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_611", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_611", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_613", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_613", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_615", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_615", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_623", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_623", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_649", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_649", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_625"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_625"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_651"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_651"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_685/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_685/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_675/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_675/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_627", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_627", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_659", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_659", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_653", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_653", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_629", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_629", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_655", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_655", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_631", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_631", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_657", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_657", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_639", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_639", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_665", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_665", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_641"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_641"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_667"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_667"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_679/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_679/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_669/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_669/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_643"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_643"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_669"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_669"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_687/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_687/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_677/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_677/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_693/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_693/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_675", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_675", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_677"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_677"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_711/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_711/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_685", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_685", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_679", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_679", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_681", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_681", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_683", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_683", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_713/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_713/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_693", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_693", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_695"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_695"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_739/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_739/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_703", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_703", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_697", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_697", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_699", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_699", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_701", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_701", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_709", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_709", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_735", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_735", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_711"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_711"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_737"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_737"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_733/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_733/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_723/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_723/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_719", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_719", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_713", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_713", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_745", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_745", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_715", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_715", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_741", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_741", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_717", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_717", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_743", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_743", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_725", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_725", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_751", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_751", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_727"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_727"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_753"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_753"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_727/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_727/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_717/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_717/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_729"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_729"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_755"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_755"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_735/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_735/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_725/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_725/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_741/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_741/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_761", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_761", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_763"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_763"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_775/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_775/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_771", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_771", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_767", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_767", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_769", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_769", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_777/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_777/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_779", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_779", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_781"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_781"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_765/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_765/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_789", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_789", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_783", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_783", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_785", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_785", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_795", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_795", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_821", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_821", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_797"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_797"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_823"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_823"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_759/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_759/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_749/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_749/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_805", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_805", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_799", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_799", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_831", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_831", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_825", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_825", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_801", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_801", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_827", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_827", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_803", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_803", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_829", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_829", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_811", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_811", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_837", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_837", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_813"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_813"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_839"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_839"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_753/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_753/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_743/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_743/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_815"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_815"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_841"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_841"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_761/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_761/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_751/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_751/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_767/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_767/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_847", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_847", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_849"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_849"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_785/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_785/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_857", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_857", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_851", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_851", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_853", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_853", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_855", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_855", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_787/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_787/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1009", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1009", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_1011"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_1011"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_181/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_181/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1019", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1019", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1013", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1013", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1015", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1015", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1017", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1017", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_185/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_185/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1025", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1025", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_4"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1027"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1027"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_175/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_175/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1185", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1185", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1035", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1035", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1029", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1029", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_1187"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_1187"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1031", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1031", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1033", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1033", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_143/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_143/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1195", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1195", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1191", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1191", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_169/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_169/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_171/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_171/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1193", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1193", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_147/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_147/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1201", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1201", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_4"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_177/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_177/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1203"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1203"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_183/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_183/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_137/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_137/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1211", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1211", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1205", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1205", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_1363"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_1363"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1207", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1207", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_105/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_105/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1371", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1371", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1367", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1367", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_131/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_131/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_133/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_133/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1369", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1369", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_109/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_109/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_139/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_139/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1379"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1379"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_145/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_145/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_99/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_99/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1387", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1387", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1383", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1383", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1385", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1385", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_93/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_93/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_95/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_95/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_101/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_101/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_107/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_107/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat/axis"], "attr": {"N": {"i": "4"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1505", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1505", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1507"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1507"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_365/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_365/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1515", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1515", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1509", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1509", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1513", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1513", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1521", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1521", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1547", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1547", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1523"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1523"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1549"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1549"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_359/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_359/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_349/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_349/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1531", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1531", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1525", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1525", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1557", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1557", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1551", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1551", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1527", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1527", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1529", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1529", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1555", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1555", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1537", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1537", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1539"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1539"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1565"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1565"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_353/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_353/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_343/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_343/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1541"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1541"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1567"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1567"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_361/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_361/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_351/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_351/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_367/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_367/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1573", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1573", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1575"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1575"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1577", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1577", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1587", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1587", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1589"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1589"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_391/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_391/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1597", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1597", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1591", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1591", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1595", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1595", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1603", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1603", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1629", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1629", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1605"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1605"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1631"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1631"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_385/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_385/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_375/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_375/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1613", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1613", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1607", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1607", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1639", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1639", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1609", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1609", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1635", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1635", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1611", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1611", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1637", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1637", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1619", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1619", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1645", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1645", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1621"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1621"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1647"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1647"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_379/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_379/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_369/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_369/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1623"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1623"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1649"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1649"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_387/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_387/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_377/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_377/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_393/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_393/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1655", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1655", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1657"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1657"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1659", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1659", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat/axis"], "attr": {"N": {"i": "4"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3065", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3065", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1781", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1781", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3067"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3067"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1783"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1783"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1487/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1487/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_289/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_289/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3075", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3075", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3069", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3069", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1791", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1791", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1785", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1785", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2383"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2383"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3071", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3071", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3073", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3073", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1789", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1789", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1749/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1749/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2391", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2391", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2385", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2385", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3081", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3081", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2387", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2387", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3083"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3083"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2389", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2389", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1797", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1797", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1823", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1823", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1753/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1753/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1481/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1481/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1799"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1799"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1825"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1825"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3091", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3091", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3085", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3085", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2397", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2397", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2399"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2399"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3087", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3087", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_283/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_283/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_273/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_273/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3089", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3089", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1807", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1807", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1801", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1801", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1833", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1833", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1827", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1827", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1743/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1743/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1803", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1803", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1829", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1829", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2401", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2401", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1805", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1805", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1831", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1831", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1475/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1475/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1477/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1477/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1813", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1813", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1839", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1839", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1483/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1483/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1815"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1815"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1841"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1841"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1737/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1737/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1739/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1739/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_277/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_277/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_267/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_267/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1745/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1745/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1817"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1817"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1843"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1843"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_285/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_285/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_275/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_275/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1751/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1751/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_291/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_291/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1849", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1849", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1851"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1851"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1853", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1853", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1863", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1863", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1865"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_1865"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_315/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_315/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1873", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1873", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1867", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1867", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1869", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1869", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1871", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1871", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1879", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1879", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1905", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1905", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1881"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1881"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1907"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1907"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_309/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_309/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_299/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_299/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1889", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1889", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1883", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1883", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1915", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1915", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1909", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1909", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1885", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1885", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1911", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1911", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1887", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1887", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_1913", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1913", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1895", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1895", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1921", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1921", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1897"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1897"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1923"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_1923"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_303/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_303/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_293/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_293/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1899"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1899"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1925"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_1925"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_311/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_311/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_301/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_301/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_317/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_317/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_1931", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1931", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1933"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_1933"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_1935", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_1935", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat/axis"], "attr": {"N": {"i": "4"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2837", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2837", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2839"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2839"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2057", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2057", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2059"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2059"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1375/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1375/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2847", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2847", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2841", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2841", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_213/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_213/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2317", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2317", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2067", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2067", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2061", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2061", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2843", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2843", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2319"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2319"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2845", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2845", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2063", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2063", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2065", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2065", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1713/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1713/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2327", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2327", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2321", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2321", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2853", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2853", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2855"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2855"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2323", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2323", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2325", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2325", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1369/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1369/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2073", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2073", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2099", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2099", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1717/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1717/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2863", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2863", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2857", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2857", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2075"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2075"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2101"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2101"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2333", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2333", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2859", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2859", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2335"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2335"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2861", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2861", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_207/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_207/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_197/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_197/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2083", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2083", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2077", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2077", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2109", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2109", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2103", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2103", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1707/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1707/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2079", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2079", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2105", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2105", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2343", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2343", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2337", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2337", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2081", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2081", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2107", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2107", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1363/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1363/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1365/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1365/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2339", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2339", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2341", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2341", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1371/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1371/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2089", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2089", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2115", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2115", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2091"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2091"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2117"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2117"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1701/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1701/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1703/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1703/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_201/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_201/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_191/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_191/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1709/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1709/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2093"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2093"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2119"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2119"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_209/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_209/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_199/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_199/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1715/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1715/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_215/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_215/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2125", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2125", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2127"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2127"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2129", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2129", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2139", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2139", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2141"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2141"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_239/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_239/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2149", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2149", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2143", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2143", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2145", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2145", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2147", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2147", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2155", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2155", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2181", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2181", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2157"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2157"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2183"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2183"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_233/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_233/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_223/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_223/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2165", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2165", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2159", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2159", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2191", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2191", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2185", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2185", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2161", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2161", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2187", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2187", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2171", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2171", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2197", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2197", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2173"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2173"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2199"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2199"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_227/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_227/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_217/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_217/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2175"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2175"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2201"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2201"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_235/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_235/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_225/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_225/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_241/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_241/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2207", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2207", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2209"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2209"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2211", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2211", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/jit_fn_/AddV2", "xs_0"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/jit_fn_/AddV2", "xs_0"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2609", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2609", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2611"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2611"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1263/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1263/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2619", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2619", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2613", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2613", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2253", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2253", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2615", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2615", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2255"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_2255"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2617", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2617", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1677/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1677/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2263", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2263", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2257", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2257", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2625", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2625", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2627"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2627"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2259", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2259", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2261", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2261", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1257/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1257/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1681/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1681/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2635", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2635", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2629", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2629", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2269", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2269", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2631", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2631", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2271"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2271"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1671/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1671/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2279", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2279", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2273", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2273", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1251/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1251/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1253/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1253/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2275", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2275", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2277", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2277", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1259/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1259/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1665/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1665/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1667/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1667/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1673/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_1673/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1679/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1679/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2413", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2413", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2415"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2415"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1183/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1183/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2423", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2423", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2417", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2417", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2419", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2419", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2421", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2421", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2455", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2455", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2431"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2431"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2457"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2457"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1177/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1177/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1167/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1167/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2439", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2439", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2433", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2433", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2465", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2465", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2459", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2459", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2435", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2435", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2461", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2461", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2437", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2437", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2463", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2463", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2445", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2445", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2471", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2471", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2447"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2447"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2473"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2473"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1171/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1171/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1161/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1161/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2449"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2449"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2475"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2475"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1179/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1179/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1169/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1169/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1185/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1185/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2481", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2481", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2483"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2483"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2485", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2485", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2495", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2495", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2497"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2497"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1209/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1209/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2505", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2505", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2499", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2499", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2501", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2501", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2503", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2503", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2537", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2537", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2513"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2513"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2539"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2539"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1203/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1203/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1193/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1193/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2521", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2521", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2515", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2515", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2547", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2547", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2541", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2541", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2517", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2517", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2543", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2543", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2519", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2519", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2545", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2545", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2527", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2527", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2529"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2529"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2555"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2555"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1197/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1197/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1187/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1187/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2531"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2531"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2557"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2557"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1205/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1205/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1195/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1195/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1211/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1211/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2565"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2565"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2567", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2567", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2577", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2577", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2579"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2579"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1261/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1261/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2587", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2587", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2581", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2581", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2583", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2583", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2585", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2585", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1265/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1265/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2595"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2595"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1247/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1247/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2603", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2603", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2597", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2597", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2599", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2599", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2601", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2601", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1241/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1241/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1243/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1243/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1249/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1249/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1267/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1267/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2641", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2641", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2643"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2643"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1295/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1295/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2651", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2651", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2645", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2645", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2647", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2647", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2649", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2649", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2657", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2657", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2683", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2683", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2659"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2659"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2685"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2685"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1289/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1289/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1279/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1279/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2667", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2667", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2661", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2661", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2693", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2693", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2687", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2687", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2663", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2663", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2689", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2689", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2665", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2665", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2691", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2691", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2673", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2673", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2699", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2699", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2675"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2675"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2701"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2701"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1283/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1283/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1273/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1273/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2677"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2677"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2703"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2703"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1291/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1291/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1281/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1281/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1297/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1297/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2709", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2709", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2711"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2711"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2713", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2713", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2723", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2723", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2725"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2725"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1321/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1321/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2733", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2733", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2727", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2727", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2729", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2729", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2731", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2731", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2741"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2741"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2767"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2767"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1315/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1315/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1305/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1305/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2749", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2749", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2743", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2743", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2775", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2775", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2769", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2769", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2745", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2745", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2771", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2771", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2747", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2747", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2773", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2773", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2755", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2755", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2781", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2781", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2757"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2757"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2783"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2783"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1309/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1309/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1299/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1299/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2759"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2759"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2785"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2785"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1317/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1317/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1307/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1307/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1323/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1323/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2791", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2791", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2793"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2793"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2795", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2795", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2805", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2805", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2807"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_2807"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1373/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1373/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2815", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2815", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2809", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2809", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2811", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2811", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2813", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2813", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1377/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1377/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2821", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2821", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2823"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2823"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1359/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1359/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2831", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2831", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2825", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2825", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2827", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2827", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2829", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2829", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1353/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1353/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1355/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1355/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1361/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1361/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1379/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1379/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2869", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2869", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2871"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2871"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1407/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1407/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2879", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2879", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2873", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2873", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2875", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2875", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2877", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2877", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2885", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2885", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2911", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2911", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2887"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2887"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2913"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2913"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1401/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1401/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1391/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1391/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2895", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2895", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2889", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2889", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2921", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2921", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2915", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2915", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2891", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2891", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2917", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2917", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2893", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2893", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2919", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2919", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2901", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2901", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2927", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2927", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2903"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2903"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2929"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2929"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1395/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1395/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1385/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1385/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2905"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2905"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2931"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2931"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1403/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1403/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1393/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1393/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1409/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1409/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2937", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2937", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2939"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2939"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2941", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2941", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2951", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2951", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2953"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_2953"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1433/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1433/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2961", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2961", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2955", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2955", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2957", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2957", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2959", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2959", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2967", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2967", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2993", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2993", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2969"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2969"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2995"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_2995"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1427/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1427/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1417/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1417/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2977", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2977", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2971", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2971", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3003", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3003", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2997", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2997", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2973", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2973", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_2999", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2999", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_2975", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2975", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3001", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3001", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_2983", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_2983", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3009", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3009", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2985"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_2985"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3011"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3011"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1421/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1421/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1411/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1411/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2987"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_2987"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3013"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3013"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1429/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1429/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1419/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1419/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1435/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1435/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3019", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3019", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3021"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3021"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3023", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3023", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3033", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3033", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3035"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3035"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1485/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1485/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3043", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3043", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3037", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3037", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3039", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3039", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3041", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3041", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1489/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_1489/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3049", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3049", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3051"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3051"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1471/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1471/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3059", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3059", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3053", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3053", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3055", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3055", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3057", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3057", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1465/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1465/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1467/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_1467/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1473/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_1473/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1491/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_1491/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6/axis"], "attr": {"N": {"i": "3"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7/axis"], "attr": {"N": {"i": "3"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8/axis"], "attr": {"N": {"i": "3"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3097", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3097", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3553", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3099"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3099"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3555"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3555"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3729", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3729", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3905", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3905", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3731"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3731"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3907"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1", "StatefulPartitionedCall/Const_3907"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1515/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1515/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_921/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_921/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3107", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3107", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3101", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3101", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3563", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3557", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3557", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_883/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_883/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_845/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_845/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3103", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3103", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3733", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3733", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3559", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3559", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3915", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3915", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3909", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3909", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3105", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3105", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3561", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3561", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3735", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3735", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3911", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3911", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3737", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3737", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3913", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3913", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3569", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3569", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3571"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3571"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3745", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3745", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3921", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3921", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3113", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3113", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3139", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3139", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3747"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3747"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3923"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3923"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3115"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3115"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3141"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3141"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_907/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_907/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3579", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3579", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3573", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3573", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_869/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_869/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_831/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_831/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1509/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1509/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1499/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1499/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3755", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3755", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3749", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3749", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3575", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3575", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3931", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3931", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3925", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3925", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3123", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3123", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3117", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3117", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3149", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3149", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3143", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3143", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3577", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3577", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3751", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3751", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3927", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3927", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3119", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3119", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3145", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3145", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3753", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3753", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3929", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3929", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3121", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3121", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3147", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3147", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_901/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_901/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_903/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_903/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_863/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_863/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_865/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_865/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_825/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_825/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_827/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_827/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3129", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3129", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3155", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3155", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3131"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3131"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3157"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3157"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_909/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_909/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_871/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_871/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_833/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat", "StatefulPartitionedCall/jax2tf_arg_833/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1503/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1503/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1493/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1493/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3133"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3133"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3159"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3159"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1511/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1511/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1501/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1501/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1517/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1517/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3165", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3165", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3167"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3167"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1551/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1551/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3175", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3175", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3171", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3171", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3173", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3173", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1553/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1553/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3183", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3183", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3185"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3185"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1541/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1541/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3193", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3193", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3187", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3187", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3191", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3191", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3199", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3199", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3225", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3225", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3201"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3201"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3227"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3227"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1535/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1535/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1525/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1525/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3203", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3203", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3235", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3235", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3229", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3229", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3205", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3205", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3231", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3231", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3207", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3207", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3233", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3233", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3215", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3215", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3241", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3241", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3217"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3217"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3243"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3243"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1529/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1529/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1519/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1519/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3219"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3219"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3245"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3245"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1537/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1537/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1527/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1527/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1543/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1543/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3251", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3251", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3253"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3253"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1561/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1561/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3261", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3261", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3255", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3255", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3257", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3257", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3259", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3259", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1563/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1563/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3269", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3269", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3271"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3271"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1589/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1589/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3279", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3279", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3273", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3273", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3275", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3275", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3277", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3277", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3285", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3285", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3311", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3311", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3287"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3287"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3313"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3313"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1583/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1583/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1573/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1573/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3295", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3295", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3289", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3289", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3321", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3321", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3315", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3315", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3291", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3291", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3317", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3317", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3293", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3293", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3319", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3319", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3301", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3301", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3327", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3327", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3303"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3303"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3329"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3329"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1577/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1577/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1567/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1567/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3305"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3305"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3331"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3331"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1585/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1585/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1575/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1575/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1591/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1591/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3337", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3337", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3339"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3339"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1625/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1625/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3347", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3347", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3341", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3341", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3343", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3343", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3345", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3345", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1627/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1627/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3355", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3355", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3357"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3357"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1615/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1615/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3359", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3359", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3363", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3363", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3371", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3371", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3397", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3397", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3373"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3373"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3399"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3399"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1609/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1609/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1599/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1599/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3375", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3375", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3401", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3401", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3379", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3379", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3387", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3387", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3413", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3413", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3389"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3389"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3415"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_3415"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1603/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1603/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1593/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1593/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3391"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3391"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3417"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_3417"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1611/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1611/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1601/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_arg_1601/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1617/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1617/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3423", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3423", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3425"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_3425"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1635/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1635/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3433", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3433", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3427", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3427", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3431", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3431", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1637/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_arg_1637/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3585", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3585", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3587"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3587"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_923/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_923/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3595", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3595", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3589", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3589", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3591", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3591", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3593", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_927/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_927/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3601", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3601", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_4"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3603"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3603"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_917/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_917/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3761", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3761", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3611", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3611", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3605", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3605", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3763"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3763"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3607", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3607", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3609", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3609", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_885/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_885/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3771", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3771", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3767", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3767", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_911/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_911/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_913/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_913/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3769", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3769", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_889/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_889/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3777", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3777", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_4"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_919/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_919/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3779"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3779"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_925/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_925/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_879/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_879/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3937", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3937", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3781", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3781", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3939"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1", "StatefulPartitionedCall/Const_3939"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3783", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3783", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3785", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3785", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_847/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_847/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3947", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3947", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3941", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3941", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3943", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3943", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit__lambda__/Tanh_1", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_873/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_873/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_875/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_875/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_3", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3945", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit__lambda__/Tanh_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3945", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit__lambda__/Tanh_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_3"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_8", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_851/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_8", "StatefulPartitionedCall/jax2tf_arg_851/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_4", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_3953", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3953", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_881/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_881/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3955"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_3955"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_887/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_887/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_841/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_841/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6"], "attr": {"N": {"i": "2"}, "equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3963", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3963", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3957", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3957", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_3959", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3959", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_3961", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_3961", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_835/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_835/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_837/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2", "StatefulPartitionedCall/jax2tf_arg_837/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_843/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat", "StatefulPartitionedCall/jax2tf_arg_843/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_9", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_849/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/Mul_9", "StatefulPartitionedCall/jax2tf_arg_849/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_5", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6"], "attr": {"equation": {"s": "Z2IsY2dlZi0+YmNlZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/jit_fn_/AddV2_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_5"], "attr": {"equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/einsum_1/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_5"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJnZCxnZi0+YWJkZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/SelectV2_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/einsum_1/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat/axis"], "attr": {"N": {"i": "4"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4081", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4081", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4083"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4083"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1107/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1107/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4091", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4091", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4085", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4085", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4087", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4087", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4089", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4089", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4097", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4097", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4123", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4123", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4099"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4099"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4125"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4125"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1101/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1101/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1091/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1091/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4107", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4107", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4101", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4101", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4133", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4133", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4127", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4127", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4103", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4103", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4129", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4129", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4105", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4105", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4131", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4131", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4113", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4113", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4139", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4139", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4115"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4115"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4141"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4141"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1095/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1095/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1085/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1085/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4117"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4117"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4143"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4143"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1103/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1103/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1093/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1093/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1109/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1109/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4149", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4149", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4151"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4151"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4153", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4153", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4163", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4165"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4165"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1133/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1133/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4173", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4173", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4167", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4167", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4169", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4171", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4171", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4179", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4179", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4205", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4205", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4181"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4181"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4207"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4207"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1127/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1127/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1117/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1117/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4189", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4183", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4183", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4215", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4215", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4209", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4185", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4185", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4211", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4211", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4187", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4187", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4213", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4213", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4195", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4195", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4221", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4221", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4197"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4197"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4223"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4223"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1121/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1121/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1111/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1111/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4199"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4199"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4225"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4225"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1129/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1129/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1119/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1119/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1135/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1135/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4231", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4231", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4233"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4233"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4235", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4235", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat/axis"], "attr": {"N": {"i": "4"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4357", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4357", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4359"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4359"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1031/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1031/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4367", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4367", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4361", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4363", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4363", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4365", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4373", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4373", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4399", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4399", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4375"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4375"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4401"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4401"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1025/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1025/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1015/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1015/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4383", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4383", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4377", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4409", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4409", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4403", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4379", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4379", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4405", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4381", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4407", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4389", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4389", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4415", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4415", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4391"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4391"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4417"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4417"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1019/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1019/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1009/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1009/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4393"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4393"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4419"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4419"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1027/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1027/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1017/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1017/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1033/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_1033/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4425", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4425", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4427"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4427"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4429", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4439", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4439", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4441"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4441"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1057/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1057/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4449", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4449", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4443", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4443", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4445", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4445", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4447", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4447", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4455", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4455", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4481", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4481", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4457"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4457"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4483"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4483"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1051/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1051/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1041/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_1041/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4465", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4465", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4459", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4459", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4491", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4491", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4485", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4485", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4461", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4461", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4487", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4487", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4463", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4463", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4489", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4489", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4471", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4471", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4497", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4497", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4473"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4473"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4499"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4499"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1045/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1045/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1035/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_1035/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4475"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4475"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4501"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4501"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1053/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1053/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1043/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_1043/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1059/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_1059/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4507", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4507", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4509"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4509"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4511", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose", "op": "Conv2DBackpropInput", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"], "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "2", "2", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat/axis"], "attr": {"N": {"i": "4"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4633", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4635"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4635"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_955/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_955/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4643", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4643", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4637", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4637", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4639", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4639", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4641", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4641", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4649", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4649", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4675", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4675", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4651"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4651"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4677"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4677"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_949/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_949/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_939/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_939/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4659", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4659", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4653", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4653", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4685", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4685", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4679", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4679", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4655", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4655", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4681", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4681", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4657", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4657", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4683", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4683", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4665", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4665", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4691", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4691", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4667"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4667"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4693"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4693"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_943/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_943/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_933/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_933/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4669"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4669"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4695"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4695"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_951/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_951/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_941/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_941/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"], "attr": {"N": {"i": "2"}, "T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_957/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat", "StatefulPartitionedCall/jax2tf_arg_957/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4701", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4701", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4703"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4703"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4705", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4705", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4715", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4715", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4717"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1", "StatefulPartitionedCall/Const_4717"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_981/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_981/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4725", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4725", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4719", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4719", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4721", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4721", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4723", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4723", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/Mul_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4731", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4731", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4757", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4757", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4733"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4733"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4759"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4759"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_975/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_975/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_965/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_arg_965/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4741", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4741", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4735", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4735", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4767", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4767", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4761", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4761", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4737", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4737", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4763", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4763", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh", "op": "Tanh", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_5"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4739", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/Const_4765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4765", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit__lambda__/Tanh"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"], "attr": {"_grappler_ArithmeticOptimizer_MinimizeBroadcasts": {"b": true}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_6", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "new_axis_mask": {"i": "0"}, "begin_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "end_mask": {"i": "0"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "op": "StridedSlice", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"], "attr": {"shrink_axis_mask": {"i": "0"}, "new_axis_mask": {"i": "0"}, "Index": {"type": "DT_INT32"}, "begin_mask": {"i": "0"}, "end_mask": {"i": "0"}, "T": {"type": "DT_FLOAT"}, "ellipsis_mask": {"i": "0"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/Mul_7", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4747", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4747", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4773", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4773", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4749"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4749"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4775"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1", "StatefulPartitionedCall/Const_4775"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_969/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_969/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_959/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose", "StatefulPartitionedCall/jax2tf_arg_959/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4751"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4751"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4777"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1", "StatefulPartitionedCall/Const_4777"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_977/ReadVariableOp"], "attr": {"N": {"i": "2"}, "equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_977/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_967/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_fn_/Mul_4", "StatefulPartitionedCall/jax2tf_arg_967/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "op": "Transpose", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"], "attr": {"Tperm": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "op": "ConcatV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"], "attr": {"N": {"i": "2"}, "Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "op": "Einsum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_983/ReadVariableOp"], "attr": {"equation": {"s": "YWJjZyxnZi0+YWJjZg=="}, "N": {"i": "2"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat", "StatefulPartitionedCall/jax2tf_arg_983/ReadVariableOp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/einsum/Einsum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1", "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1", "op": "Square", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub", "op": "Sub", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/ArithmeticOptimizer/ReplaceMulWithSquare_Mul_1"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/Const_4783", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4783", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4785"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1", "StatefulPartitionedCall/Const_4785"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "op": "Rsqrt", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Rsqrt", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__lambda__/Sub_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_4"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/Mul_5", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit_fn_/AddV2_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "op": "GreaterEqual", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/Const_4787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/Const_4787", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "op": "SelectV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/GreaterEqual", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Mul"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit__where_/SelectV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "op": "Sum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"], "attr": {"T": {"type": "DT_FLOAT"}, "Tidx": {"type": "DT_INT32"}, "keep_dims": {"b": false}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "op": "Reshape", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"], "attr": {"T": {"type": "DT_FLOAT"}, "Tshape": {"type": "DT_INT32"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "op": "BroadcastTo", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"], "attr": {"Tidx": {"type": "DT_INT32"}, "T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "op": "RealDiv", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/truediv", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "op": "Maximum", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "VkFMSUQ="}, "data_format": {"s": "TkhXQw=="}, "use_cudnn_on_gpu": {"b": true}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Maximum", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg", "op": "Neg", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp", "op": "Exp", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Neg"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Exp"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv", "op": "Reciprocal", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "op": "Mul", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/truediv"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_fn_/Mul", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/jit_fn_/AddV2_2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/jit_fn_/AddV2"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Conv2D", "op": "Conv2D", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/transpose_3"], "device": "/device:CPU:0", "attr": {"padding": {"s": "U0FNRQ=="}, "data_format": {"s": "TkhXQw=="}, "strides": {"list": {"i": ["1", "1", "1", "1"]}}, "use_cudnn_on_gpu": {"b": true}, "T": {"type": "DT_FLOAT"}, "explicit_paddings": {"list": {}}, "dilations": {"list": {"i": ["1", "1", "1", "1"]}}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/jit_fn_/AddV2", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/transpose_3"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/jit_fn_/AddV2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Reshape"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Conv2D", "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Reshape"]}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit_fn_/AddV2_2", "op": "AddV2", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/jit_fn_/AddV2", "xs_0"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/jit_fn_/AddV2", "xs_0"]}, {"name": "Identity", "op": "Identity", "input": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit_fn_/AddV2_2"], "attr": {"T": {"type": "DT_FLOAT"}}, "inputs": ["StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit_fn_/AddV2_2"]}], "library": {}, "versions": {"producer": 1395}}, "weightsManifest": [{"paths": ["group1-shard1of14.bin", "group1-shard2of14.bin", "group1-shard3of14.bin", "group1-shard4of14.bin", "group1-shard5of14.bin", "group1-shard6of14.bin", "group1-shard7of14.bin", "group1-shard8of14.bin", "group1-shard9of14.bin", "group1-shard10of14.bin", "group1-shard11of14.bin", "group1-shard12of14.bin", "group1-shard13of14.bin", "group1-shard14of14.bin"], "weights": [{"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4787", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4783", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4785", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4747", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4741", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4739", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4737", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4735", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4731", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4733", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_975/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4749", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_969/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4751", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_977/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4773", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4767", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4765", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4763", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4761", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4757", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4725", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4723", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4721", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4719", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4715", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4717", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_981/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4759", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_965/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4775", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_959/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4777", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_967/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_983/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4705", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4701", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4703", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4665", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4659", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4657", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4655", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4653", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4649", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4651", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_949/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4667", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_943/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4669", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_951/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4691", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4685", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4683", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4681", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4679", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4675", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4643", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4641", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4639", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4637", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4633", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4635", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_955/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4677", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_939/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4693", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_933/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4695", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_941/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_957/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/transpose_4", "shape": [2, 2, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4511", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4507", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4509", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4471", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4465", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4463", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4461", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4459", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4455", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4457", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1051/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4473", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1045/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4475", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1053/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4497", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4491", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4489", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4487", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4485", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4481", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4449", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4447", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4445", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4443", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4439", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4441", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1057/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4483", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1041/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4499", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1035/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4501", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1043/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1059/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4429", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4425", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4427", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4389", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4383", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4381", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4379", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4377", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4373", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4375", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1025/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4391", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1019/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4393", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1027/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4415", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4409", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4407", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4405", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4403", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4399", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4367", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4365", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4363", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4361", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4357", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4359", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1031/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4401", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1015/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4417", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1009/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4419", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1017/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1033/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/transpose_4", "shape": [2, 2, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4235", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4231", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4233", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4195", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4189", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4187", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4185", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4183", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4179", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4181", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1127/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4197", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1121/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4199", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1129/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4221", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4215", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4213", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4211", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4209", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4205", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4173", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4171", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4169", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4167", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4163", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4165", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1133/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4207", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1117/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4223", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1111/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4225", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1119/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1135/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4153", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4149", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4151", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4113", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4107", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4105", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4103", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4101", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4097", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4099", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1101/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4115", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1095/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4117", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1103/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4139", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4133", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4131", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4129", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4127", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4123", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4091", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4089", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4087", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4085", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4081", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4083", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1107/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4125", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1091/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4141", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1085/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4143", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1093/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1109/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/transpose_4", "shape": [2, 2, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_27/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3989", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3987", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3979", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3969", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3973", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3967", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4017", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4015", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4007", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3997", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4001", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3995", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/transpose_3", "shape": [1, 1, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4045", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4043", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4035", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4025", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4029", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4023", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4073", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4071", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4063", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4053", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4057", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4051", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3", "shape": [1, 1, 512, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4265", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4263", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4255", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4245", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4249", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4243", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4293", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4291", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4283", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4273", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4277", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4271", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/transpose_3", "shape": [1, 1, 128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_31/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4321", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4319", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4311", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4301", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4305", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4299", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4349", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4347", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4339", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4329", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4333", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4327", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/transpose_3", "shape": [1, 1, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3", "shape": [1, 1, 256, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4541", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4539", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4531", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4521", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4525", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4519", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_911/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3611", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3609", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3607", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3605", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3601", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3603", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_917/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_913/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_919/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_925/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4569", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4567", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4559", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4549", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4553", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4547", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4597", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4595", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4587", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4577", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4581", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4575", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_873/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3787", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3785", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3783", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3781", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3777", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3779", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_879/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_875/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_881/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_887/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4625", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4623", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4615", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4605", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_4609", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_4603", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/transpose_3", "shape": [1, 1, 64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_835/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3963", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3961", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3959", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3957", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3953", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3955", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_841/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_837/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_843/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_849/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3947", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3945", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3943", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3941", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3937", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3939", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_847/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_825/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3931", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3929", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3927", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3925", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3921", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3915", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3913", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3911", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3909", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3905", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3907", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_845/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/in_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3923", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_831/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_827/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_833/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_851/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/out_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/transpose_4", "shape": [2, 2, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3771", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3769", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3767", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3765", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3761", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3763", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_885/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_863/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3755", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3753", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3751", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3749", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3745", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3739", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3737", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3735", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3733", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3729", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3637", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3635", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3627", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3617", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3621", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3615", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3665", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3663", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3655", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3645", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3649", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3643", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/transpose_3", "shape": [1, 1, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_22/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3693", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3691", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3683", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3673", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3677", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3671", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3721", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3719", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3711", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3701", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3705", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3699", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/transpose_3", "shape": [1, 1, 128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_7/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/transpose_3", "shape": [1, 1, 192, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3731", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_883/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/in_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3747", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_869/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_865/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_871/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_889/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/out_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/transpose_4", "shape": [2, 2, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3595", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3593", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3591", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3589", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3585", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3587", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_923/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_901/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3579", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3577", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3575", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3573", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3569", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3563", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3561", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3559", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3557", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3553", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3461", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3459", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3451", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3441", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3445", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3439", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3489", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3487", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3479", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3469", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3473", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3467", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3517", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3515", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3507", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3497", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3501", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3495", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3545", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3543", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3535", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3525", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3529", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3523", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/transpose_3", "shape": [1, 1, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_20/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_6/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/transpose_3", "shape": [1, 1, 384, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3555", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_921/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/in_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3571", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_907/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_903/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_909/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_927/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/out_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/transpose_4", "shape": [2, 2, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3433", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3431", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3429", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3427", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3423", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3387", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3381", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3379", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3377", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3375", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3371", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3373", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1609/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3389", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1603/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3391", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1611/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3413", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3407", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3405", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3403", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3401", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3397", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3365", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3363", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3361", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3359", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3355", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3357", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1615/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3399", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1599/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3415", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1593/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3417", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1601/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1617/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3347", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3345", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3343", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3341", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3337", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3301", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3295", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3293", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3291", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3289", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3285", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3287", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1583/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3303", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1577/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3305", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1585/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3327", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3321", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3319", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3317", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3315", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3311", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3279", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3277", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3275", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3273", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3269", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3271", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1589/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3313", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1573/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3329", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1567/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3331", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1575/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1591/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3339", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1625/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1627/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3425", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1635/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1637/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3261", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3259", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3257", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3255", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3251", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3215", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3209", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3207", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3205", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3203", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3199", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3201", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1535/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3217", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1529/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3219", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1537/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3241", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3235", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3233", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3231", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3229", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3225", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3193", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3191", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3189", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3187", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3183", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3185", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1541/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3227", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1525/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3243", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1519/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3245", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1527/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1543/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3175", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3173", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3171", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3169", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3165", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3129", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3123", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3121", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3119", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3117", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3113", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3115", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1509/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3131", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1503/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3133", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1511/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3155", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3149", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3147", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3145", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3143", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3139", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3107", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3105", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3103", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3101", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3097", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3099", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1515/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3141", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1499/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3157", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1493/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3159", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1501/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1517/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3167", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1551/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1553/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3253", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1561/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1563/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/transpose_3", "shape": [4, 4, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_0/input_proj/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_global_block_1/input_proj/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_2/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_1/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_24/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3813", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3811", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3803", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3793", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3797", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3791", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3841", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3839", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3831", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3821", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3825", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3819", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/transpose_3", "shape": [1, 1, 64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3869", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3867", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3859", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3849", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3853", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3847", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1475/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3091", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3089", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3087", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3085", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3081", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3083", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1481/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1477/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1483/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1489/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3075", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3073", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3071", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3069", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3065", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3067", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1487/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1465/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3059", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3057", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3055", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3053", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3049", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3043", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3041", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3039", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3037", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3033", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3035", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1485/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/in_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3051", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1471/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1467/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1473/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1491/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/out_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3023", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3019", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3021", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2983", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2977", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2975", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2973", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2971", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2967", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2969", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1427/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2985", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1421/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2987", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1429/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3009", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3003", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3001", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2999", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2997", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2993", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2961", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2959", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2957", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2955", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2951", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2953", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1433/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2995", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1417/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3011", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1411/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3013", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1419/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1435/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2941", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2937", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2939", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2901", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2895", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2893", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2891", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2889", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2885", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2887", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1401/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2903", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1395/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2905", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1403/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2927", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2921", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2919", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2917", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2915", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2911", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2879", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2877", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2875", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2873", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2869", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2871", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1407/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2913", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1391/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2929", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1385/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2931", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1393/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1409/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1363/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2863", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2861", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2859", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2857", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2853", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2855", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1369/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1365/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1371/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1377/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2847", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2845", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2843", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2841", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2837", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2839", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1375/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1353/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2831", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2829", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2827", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2825", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2821", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2815", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2813", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2811", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2809", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2805", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2807", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1373/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/in_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2823", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1359/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1355/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1361/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1379/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/out_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2795", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2791", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2793", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2755", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2749", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2747", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2745", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2743", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2739", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2741", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1315/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2757", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1309/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2759", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1317/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2781", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2775", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2773", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2771", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2769", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2765", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2733", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2731", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2729", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2727", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2723", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2725", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1321/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2767", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1305/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2783", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1299/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2785", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1307/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1323/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2713", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2709", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2711", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2673", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2667", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2665", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2663", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2661", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2657", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2659", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1289/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2675", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1283/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2677", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1291/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2699", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2693", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2691", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2689", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2687", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2683", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2651", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2649", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2647", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2645", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2641", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2643", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1295/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2685", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1279/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2701", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1273/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2703", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1281/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1297/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1251/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2635", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2633", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2631", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2629", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2625", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2627", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1257/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1253/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1259/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1265/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2619", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2617", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2615", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2613", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2609", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2611", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1263/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1241/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2603", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2601", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2599", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2597", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2593", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2587", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2585", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2583", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2581", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2577", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2579", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1261/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/in_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2595", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1247/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1243/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1249/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1267/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/out_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2567", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2563", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2565", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2527", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2521", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2519", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2517", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2515", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2511", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2513", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1203/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2529", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1197/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2531", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1205/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2553", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2547", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2545", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2543", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2541", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2537", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2505", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2503", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2501", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2499", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2495", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2497", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1209/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2539", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1193/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2555", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1187/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2557", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1195/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1211/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2485", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2481", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2483", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2445", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2439", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2437", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2435", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2433", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2429", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2431", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1177/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2447", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1171/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2449", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1179/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2471", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2465", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2463", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2461", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2459", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2455", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2423", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2421", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2419", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2417", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2413", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2415", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1183/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2457", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1167/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2473", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1161/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2475", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1169/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1185/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1665/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2279", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2277", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2275", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2273", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2269", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2271", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1671/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1667/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1673/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1679/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2263", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2261", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2259", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2257", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2253", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2255", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1677/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1655/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2247", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2245", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2243", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2241", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2237", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2231", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2229", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2227", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2225", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2221", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2223", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1675/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/in_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2239", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1661/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1657/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1663/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1681/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/out_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/transpose_3", "shape": [3, 3, 32, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_1/Reshape", "shape": [1, 1, 1, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/transpose_3", "shape": [3, 3, 3, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_0/Conv_2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2211", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2207", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2209", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2171", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2165", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2163", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2161", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2159", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2155", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2157", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_233/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2173", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_227/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2175", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_235/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2197", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2191", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2189", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2187", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2185", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2181", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2149", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2147", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2145", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2143", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2139", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2141", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_239/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2183", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_223/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2199", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_217/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2201", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_225/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_241/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2129", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2125", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2127", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2089", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2083", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2081", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2079", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2077", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2073", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2075", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_207/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2091", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_201/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2093", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_209/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2115", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2109", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2107", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2105", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2103", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2099", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2067", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2065", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2063", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2061", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2057", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2059", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_213/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2101", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_197/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2117", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_191/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2119", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_199/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_215/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/transpose_4", "shape": [2, 2, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1965", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1963", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1955", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1945", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1949", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1943", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1993", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1991", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1983", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1973", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1977", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1971", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2021", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2019", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2011", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2001", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2005", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1999", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2049", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2047", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2039", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2029", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2033", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2027", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/transpose_3", "shape": [1, 1, 64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_17/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/transpose_3", "shape": [3, 3, 3, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/transpose_3", "shape": [1, 1, 64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/cross_gating_block/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/transpose_3", "shape": [4, 4, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_0/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1701/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2343", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2341", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2339", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2337", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2333", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2335", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1707/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1703/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1709/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1715/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2327", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2325", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2323", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2321", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2317", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2319", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1713/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1691/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2311", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2309", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2307", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2305", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2301", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2295", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2293", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2291", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2289", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2285", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2287", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1711/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/in_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2303", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1697/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1693/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1699/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1717/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/out_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/transpose_3", "shape": [3, 3, 64, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_1/Reshape", "shape": [1, 1, 1, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/transpose_3", "shape": [3, 3, 3, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_1/Conv_2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1935", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1931", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1933", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1895", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1889", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1887", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1885", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1883", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1879", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1881", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_309/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1897", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_303/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1899", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_311/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1921", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1915", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1913", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1911", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1909", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1905", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1873", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1871", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1869", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1867", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1863", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1865", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_315/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1907", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_299/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1923", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_293/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1925", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_301/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_317/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1853", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1849", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1851", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1813", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1807", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1805", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1803", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1801", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1797", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1799", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_283/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1815", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_277/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1817", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_285/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1839", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1833", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1831", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1829", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1827", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1823", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1791", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1789", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1787", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1785", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1781", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1783", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_289/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1825", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_273/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1841", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_267/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1843", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_275/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_291/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/transpose_4", "shape": [2, 2, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1689", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1687", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1679", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1669", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1673", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1667", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1717", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1715", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1707", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1697", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1701", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1695", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/transpose_3", "shape": [1, 1, 128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_13/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1745", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1743", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1735", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1725", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1729", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1723", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1773", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1771", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1763", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1753", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1757", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1751", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/transpose_3", "shape": [1, 1, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/transpose_3", "shape": [1, 1, 256, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_1/UNetEncoderBlock_0/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/transpose_3", "shape": [3, 3, 3, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_1/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/transpose_3", "shape": [1, 1, 96, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/cross_gating_block/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/transpose_3", "shape": [4, 4, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_1/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1737/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2407", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2405", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2403", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2401", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2397", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2399", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1743/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1739/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1745/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1751/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2391", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2389", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2387", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2385", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2381", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2383", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1749/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1727/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2375", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2373", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2371", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2369", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2365", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2359", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2357", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2355", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_2353", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2349", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2351", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1747/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/in_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_2367", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1733/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1729/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1735/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_1753/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/out_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/transpose_3", "shape": [3, 3, 128, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_1/Reshape", "shape": [1, 1, 1, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/transpose_3", "shape": [3, 3, 3, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_supervised_attention_module_2/Conv_2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1659", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1655", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1657", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1619", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1613", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1611", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1609", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1607", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1603", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1605", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_385/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1621", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_379/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1623", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_387/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1645", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1639", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1637", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1635", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1633", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1629", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1597", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1595", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1593", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1591", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1587", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1589", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_391/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1631", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_375/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1647", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_369/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1649", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_377/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_393/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1577", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1573", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1575", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1537", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1531", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1529", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1527", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1525", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1521", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1523", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_359/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1539", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_353/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1541", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_361/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1563", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1557", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1555", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1553", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1551", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1547", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1515", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1513", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1511", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1509", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1505", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1507", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_365/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1549", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_349/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1565", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_343/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1567", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_351/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_367/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/transpose_4", "shape": [2, 2, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_169/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1035", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1033", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1031", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1029", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1025", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1027", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_175/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_171/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_177/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_183/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_9/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1413", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1411", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1403", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1393", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1397", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1391", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_131/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1211", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1209", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1207", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1205", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1201", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1203", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_137/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_133/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_139/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_145/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1441", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1439", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1431", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1421", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1425", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1419", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/transpose_3", "shape": [1, 1, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1469", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1467", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1459", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1449", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1453", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1447", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_93/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1387", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1385", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1383", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1381", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1377", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1379", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_99/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_95/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_101/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_y/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_107/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1371", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1369", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1367", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1365", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1361", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1363", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_105/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_83/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1355", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1353", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1351", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1349", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1345", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1339", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1337", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1335", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1333", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1329", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1331", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_103/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/in_project_x/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1347", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_89/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_85/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_91/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_109/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/out_project_y/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/transpose_4", "shape": [2, 2, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1195", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1193", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1191", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1189", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1185", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1187", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_143/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_121/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1179", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1177", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1175", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1173", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1169", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1163", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1161", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1159", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1157", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1153", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1061", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1059", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1051", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1041", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1045", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1039", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1089", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1087", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1079", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1069", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1073", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1067", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/transpose_3", "shape": [1, 1, 32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_4/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1117", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1115", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1107", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1097", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1101", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1095", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1145", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1143", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1135", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1125", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1129", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1123", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/transpose_3", "shape": [1, 1, 128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_1/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/transpose_3", "shape": [1, 1, 192, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1155", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_141/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/in_project_x/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1171", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_127/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_123/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_129/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_147/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/out_project_y/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/transpose_4", "shape": [2, 2, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1019", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1017", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1015", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1013", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1009", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1011", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_y/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_181/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_159/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1003", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1001", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_999", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_997", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_993", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_987", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_985", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_983", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_981", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_977", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_885", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_883", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_875", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_865", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_869", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_863", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_913", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_911", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_903", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_893", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_897", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_891", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_941", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_939", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_931", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_921", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_925", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_919", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_969", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_967", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_959", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_949", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_953", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_947", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/transpose_3", "shape": [1, 1, 64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/transpose_3", "shape": [1, 1, 384, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_979", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/LayerNorm_x/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_179/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/in_project_x/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_995", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_165/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_4/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_7/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_2/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_161/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Dense_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/jit_swapaxes_/transpose_3/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_6/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/transpose_10/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/Reshape_7/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_167/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/SplitHeadMultiAxisGating_x/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_185/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/out_project_y/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/conv2d_transpose/input_sizes/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/transpose_4", "shape": [2, 2, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_857", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_855", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_853", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_851", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_847", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_811", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_805", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_803", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_801", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_799", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_795", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_797", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_759/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_813", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_753/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_815", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_761/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_837", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_831", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_829", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_827", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_825", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_821", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_789", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_787", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_785", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_783", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_779", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_781", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_765/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_823", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_749/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_839", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_743/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_841", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_751/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_767/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_771", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_769", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_767", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_765", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_761", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_725", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_719", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_717", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_715", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_713", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_709", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_711", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_733/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_727", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_727/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_729", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_735/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_751", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_745", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_743", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_741", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_739", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_735", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_703", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_701", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_699", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_697", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_693", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_695", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_739/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_737", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_723/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_753", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_717/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_755", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_725/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_741/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_763", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_775/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_777/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_849", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_785/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_787/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_685", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_683", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_681", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_679", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_675", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_639", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_633", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_631", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_629", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_627", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_623", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_625", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_685/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_641", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_679/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_643", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_687/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_665", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_659", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_657", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_655", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_653", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_649", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_617", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_615", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_613", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_611", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_607", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_609", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_691/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_651", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_675/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_667", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_669/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_669", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_677/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_693/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_599", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_597", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_595", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_593", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_589", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_553", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_547", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_545", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_543", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_541", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_537", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_539", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_659/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_555", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_653/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_557", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_661/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_579", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_573", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_571", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_569", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_567", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_563", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_531", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_529", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_527", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_525", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_521", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_523", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_665/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_565", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_649/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_581", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_643/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_583", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_651/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_667/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_591", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_701/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_703/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_0/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_677", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_711/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_713/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_mixing/Dense_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/channel_attention_block_1_1/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/transpose_3", "shape": [4, 4, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_0/input_proj/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_global_block_1/input_proj/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_2/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/transpose_3", "shape": [1, 1, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_1/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/ConvTranspose_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_6/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1237", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1235", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1227", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1217", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1221", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1215", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1265", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1263", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1255", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1245", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1249", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1243", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/transpose_3", "shape": [1, 1, 64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1293", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1291", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_4/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1283", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_2", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1273", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1277", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_21", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_3/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1271", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_3", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_6", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_511", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_507", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_509", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_471", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_465", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_463", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_461", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_459", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_455", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_457", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_613/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_473", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_607/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_475", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_615/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_497", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_491", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_489", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_487", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_485", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_481", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_449", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_447", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_445", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_443", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_439", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_441", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_619/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_483", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_603/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_499", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_597/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_501", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_605/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_621/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_429", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_425", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_427", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_389", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_383", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_381", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_379", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_377", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_373", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_375", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_587/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_391", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_581/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_393", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_589/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_415", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_409", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_407", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_405", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_403", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_399", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_367", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_365", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_363", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_361", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_357", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_359", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_593/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_401", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_577/ReadVariableOp", "shape": [128, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_417", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_571/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_419", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_579/ReadVariableOp", "shape": [128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_595/ReadVariableOp", "shape": [256, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_347", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_343", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_345", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_307", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_301", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_299", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_297", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_295", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_291", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_293", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_537/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_309", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_531/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_311", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_539/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_333", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_327", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_325", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_323", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_321", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_317", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_285", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_283", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_281", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_279", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_275", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_277", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_543/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_319", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_527/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_335", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_521/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_337", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_529/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_545/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_265", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_261", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_263", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 64, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 16], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 16, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_225", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_219", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_217", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_215", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_213", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_209", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_211", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_511/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_227", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_505/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_229", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_513/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_251", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_245", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_243", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_241", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_239", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_235", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_203", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_201", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_199", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_197", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_193", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_195", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_517/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_237", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_501/ReadVariableOp", "shape": [64, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_253", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_495/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_255", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_503/ReadVariableOp", "shape": [64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_519/ReadVariableOp", "shape": [128, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_183", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_179", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_181", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_11/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_143", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_137", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_135", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_133", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_131", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_127", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_129", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_461/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_145", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_455/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_147", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_463/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_169", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_163", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_161", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_159", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_157", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_153", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_121", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_119", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_117", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_115", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_111", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_113", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_467/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_155", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_451/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_171", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_445/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_173", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_453/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_469/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_1/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_sigmoid_/Const_1", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/jit_leaky_relu_/jit_fn_/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_101", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_97", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_99", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/transpose_3", "shape": [3, 3, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/conv2/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [2], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/jit__reduce_sum_/BroadcastTo/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit__mean_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/transpose_3", "shape": [1, 1, 32, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_0/Reshape", "shape": [1, 1, 1, 8], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/jit_relu_/jit__lambda__/Cast", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/transpose_3", "shape": [1, 1, 8, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/channel_attention_block_10/channel_attention/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_61", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_55", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_53", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_51", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_49", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_45", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_47", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_435/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_63", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_429/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/GridGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_65", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_437/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/GridGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_87", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_81", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_79", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_77", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_75", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_71", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_39", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_37", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_35", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_33", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_29", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_1_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_31", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/LayerNorm_in/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_441/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_1/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_73", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/LayerNorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_425/ReadVariableOp", "shape": [32, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/in_project/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_1/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/strided_slice_1/stack_2", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/jit__reduce_sum_/Sum/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/jit__mean_/truediv_recip", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/2", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_1/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_89", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_2", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/intermediate_layernorm/Reshape_3", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_419/ReadVariableOp", "shape": [256, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/Dense_0/Reshape", "shape": [1, 1, 1, 256], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/BlockGatingUnit/jit_swapaxes_/transpose_1/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_91", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_427/ReadVariableOp", "shape": [32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_2/shape/5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/transpose_4/perm", "shape": [6], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_27", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/BlockGmlpLayer/Reshape_3/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_arg_443/ReadVariableOp", "shape": [64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/SplitHeadMultiAxisGmlpLayer_0/out_project/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/transpose_3", "shape": [3, 3, 3, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/transpose_3", "shape": [1, 1, 64, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/transpose_3", "shape": [4, 4, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_0/Conv_1/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_11", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_9", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_1/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_17", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/y", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_25", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_29", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_31", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_2/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_15", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_1/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_35", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_13", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_37", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_4/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_4/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_53", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/y", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_1/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_1/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_5/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/transpose_3", "shape": [3, 3, 3, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/transpose_3", "shape": [1, 1, 96, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_0/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/transpose_3", "shape": [4, 4, 64, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_1/Conv_1/Reshape", "shape": [1, 1, 1, 64], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/transpose_3", "shape": [3, 3, 3, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_input_conv_2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/transpose_3", "shape": [1, 1, 192, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_encoder_block_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1321", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1319", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1311", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1301", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1305", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1299", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_2/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/transpose_3", "shape": [1, 1, 96, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_cross_gating_block_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1497", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1495", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1487", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1477", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_1481", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1475", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/transpose_3", "shape": [1, 1, 32, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/transpose_3", "shape": [1, 1, 512, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_0_decoder_block_2/UNetEncoderBlock_0/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_1/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_19", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_75", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_69", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_17", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_71", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_73", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_77", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_7/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_81", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_7/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_85", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/y", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_2/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_2/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_91", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_93", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_8/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_97", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_99", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_8/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_23", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_109", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_103", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_21", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_105", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_107", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_111", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_10/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_115", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_10/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_121", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/y", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/clip_by_value_3/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/GatherV2_3/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_125", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_127", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Reshape_11/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_1", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice/stack_2", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Sign_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_1/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/add/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_1", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_1/stack_2", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_131", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Const_133", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_11/shape/3", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/transpose_3", "shape": [3, 3, 3, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_conv_2/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_input_fuse_sam_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/transpose_3", "shape": [1, 1, 192, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/transpose_3", "shape": [1, 1, 128, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_encoder_block_2/cross_gating_block/Conv_0/Reshape", "shape": [1, 1, 1, 128], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3897", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3895", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_9/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_65", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_57", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_59", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3887", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_43", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_6", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_5", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__lambda__/Cast_4", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_2/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_39", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3877", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/Const_3881", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_41", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_6/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_45", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/start", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/range_3/delta", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_51", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Reshape_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_55", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_8/shape/1", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/Const_7", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_3875", "shape": [], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Sum_1/reduction_indices", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Reshape_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_9", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/BroadcastTo_1/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__reduce_sum_/Const_13", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_7", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/BroadcastTo_5/shape/0", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_33", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_10", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_47", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_49", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Reshape_13", "shape": [1, 1], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_61", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Sign_4", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/NotEqual_3/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/add_1/y", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_1", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/strided_slice_2/stack_2", "shape": [1], "dtype": "int32"}, {"name": "StatefulPartitionedCall/Const_5", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__where_/jit__broadcast_arrays_/Const_63", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/jit__einsum_/transpose/perm", "shape": [4], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/concat_8/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/transpose_3", "shape": [1, 1, 96, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_cross_gating_block_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/transpose_3", "shape": [1, 1, 32, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_35/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/concat/axis", "shape": [], "dtype": "int32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/transpose_3", "shape": [1, 1, 128, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_decoder_block_0/UNetEncoderBlock_0/Conv_0/Reshape", "shape": [1, 1, 1, 32], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/transpose_3", "shape": [3, 3, 32, 3], "dtype": "float32"}, {"name": "StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/stage_1_output_conv_0/Reshape", "shape": [1, 1, 1, 3], "dtype": "float32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_2/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub/x_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_6/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_2/x_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_28/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_29/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_30/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_32/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_33/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_34/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_21/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_23/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_18/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_19/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_25/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_15/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_16/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_12/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_14/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_10/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_3/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_5/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_0/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_1/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_7/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_3/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_1/x_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_8/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_11/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/BroadcastTo_9/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/jit__resize_/Sub_3/x_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_5/shape_const_axis", "shape": [], "dtype": "int32"}, {"name": "ConstantFolding/StatefulPartitionedCall/jax2tf_apply_/Wrapper/backbone/UpSampleRatio_26/jit__resize_/BroadcastTo_7/shape_const_axis", "shape": [], "dtype": "int32"}]}]}